package main

import (
	"fmt"
	"github.com/mitchellh/mapstructure"
	"github.com/spf13/viper"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
	"log"
	"os"
	"strconv"
)

const _defaultConf = "app.yaml"

type SshAddr struct {
	IP string `mapstructure:"IP"`
	User string `mapstructure:"User"`
	Pwd string `mapstructure:"Pwd"`
}

func Viper(path ...string) (*viper.Viper, error) {
	v := viper.New()

	if len(path) == 0 {
		v.SetConfigFile(_defaultConf)
	} else {
		v.SetConfigFile(path[0])
	}
	err := v.ReadInConfig()
	if err != nil {
		return nil, err
	}

	return v, nil
}

// GetSshAp 获取账号和密码列表
func GetSshAp() ([]*SshAddr, error) {
	vi, err := Viper()
	if err != nil {
		log.Println(err)
		return nil, err
	}
	data := vi.Get("addr").([]interface{})
	res := make([]*SshAddr, 0)
	var mr SshAddr
	for _, v := range data {
		mapstructure.Decode(v, &mr)
		res = append(res, &SshAddr{
			IP: mr.IP,
			User: mr.User,
			Pwd: mr.Pwd,
		})
	}

	return res, nil
}

func sshConn(user string, pwd string, ip string)  {
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pwd),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	client, err := ssh.Dial("tcp", ip, config)
	if err != nil {
		log.Fatal("Failed to dial: ", err)
	}

	session, err := client.NewSession()

	if err != nil {
		panic(err)
	}

	defer session.Close()
	fd := int(os.Stdin.Fd())

	state, err := terminal.MakeRaw(fd)
	if err != nil {
		panic(err)
	}

	defer terminal.Restore(fd, state)

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	// Request pseudo terminal
	if err := session.RequestPty("xterm", 40, 80, modes); err != nil {
		log.Fatal("request for pseudo terminal failed: ", err)
	}

	if err := session.Shell(); err != nil {
		log.Fatal("failed to start shell: ", err)
	}

	fmt.Println(session.Wait())
}

// 简单启动就是查询, 加上 id 启动直接链接即可
func main() {
	fmt.Println("输入列表id 即可直接链接")

	data, err := GetSshAp()
	if err != nil {
		return
	}

	for k, v := range data {
		fmt.Println(k," -> ", v.IP)
	}

	if len(os.Args) == 1 {
		return
	}

	sid := os.Args[1]
	id, err := strconv.Atoi(sid)
	if err != nil {
		return
	}

	for k, v := range data {
		if k == id {
			fmt.Println("开始链接", v.User, v.Pwd, v.IP)
			sshConn(v.User, v.Pwd, v.IP)
		}
	}
}
